#ifndef _FEMSPACE_H_
#define _FEMSPACE_H_

#include "Dofs.h"
#include "Mesh.h"
#include "Point.h"
#include "TemplateElement.h"
#include "AMG.cpp"
#include "Equation.h"
#include <Eigen/Sparse>
#include <Eigen/IterativeLinearSolvers>
#include <Eigen/Dense>
#include <ctime>
#include <cmath>

/**
 * @file   FEMSpace.h
 * @author AOM
 * @date   Sat Mar  4 10:24:15 2023
 * 
 * @brief  有限元空间基类
 *         主要功能：生成网格，分配自由度，组装刚度矩阵、右端项，处理边界条件，求解。
 * 
 * 
 */

typedef Eigen::SparseMatrix<double> SpMat;
typedef Eigen::VectorXd RHS;
typedef Eigen::VectorXd Solution;

template <unsigned int DIM>
std::map<int,double> interpolateBoundaryValues(auto& _boundary_function, vector<Dofs<DIM>>& _boundary_dofs);

void applyBoundaryValues(std::map<int,double>& boundary_values,Eigen::SparseMatrix<double>& system_matrix,Eigen::VectorXd& system_rhs);

template <unsigned int DIM>
class FEMSpace
{
public:
    //virtual void make_grid() = 0;
    virtual void setup_system() = 0;
    virtual void assemble_system() = 0;
    virtual void dealwithBoundaryCondition() = 0;
    virtual void solve() = 0;
};

class Possion_2D : public FEMSpace<2>{
public:
    Possion_2D();
    Possion_2D(Mesh<2>*,TemplateElement<2>*,Equation<2>*);
    Possion_2D(const Possion_2D& );
    ~Possion_2D(){/*delete _mesh;delete _element;delete _eqation;*/};
    //virtual void make_grid();
    virtual void setup_system();
    virtual void assemble_system();
    virtual void dealwithBoundaryCondition();
    virtual void solve();

    Solution getSolution() const;
    SpMat A() const;
    RHS Rhs() const;
protected:
    SpMat system_matrix;
    RHS system_rhs;
    Solution solution;
    //Rectangle2D* _domain;
    Mesh<2>* _mesh;
    TemplateElement<2>* _element;
    Equation<2>* _eqation;
    //BoundaryCondition<2> _boundarycondition;
};

Solution Possion_2D::getSolution() const
{
    return solution;
}

Possion_2D::Possion_2D()
{   
    system_matrix = SpMat(int(_mesh->getTotalNumDofs()), int(_mesh->getTotalNumDofs()));
    system_rhs = RHS(int(_mesh->getTotalNumDofs()));
    //_boundarycondition.apply(_mesh);
}
Possion_2D::Possion_2D(Mesh<2>* mesh,TemplateElement<2>* element,Equation<2>* equation)
{
    //_domain = domain;
    _mesh = mesh;
    _element = element;
    _eqation = equation;
    //_boundarycondition = (BoundaryCondition);
    system_matrix = SpMat(int(_mesh->getTotalNumDofs()), int(_mesh->getTotalNumDofs()));
    system_rhs = RHS(int(_mesh->getTotalNumDofs()));
    //_boundarycondition.apply(_mesh);
}

Possion_2D::Possion_2D(const Possion_2D& p)
{
    //_domain = p._domain;
    _mesh = p._mesh;
    _element = p._element;
    _eqation = p._eqation;
    //_boundarycondition = p._boundarycondition;
    system_matrix = p.system_matrix;
    system_rhs = p.system_rhs;
    //_boundarycondition.apply(_mesh);
}

// void Possion_2D::make_grid(){
//     vector<Point<2>> p = {{-1.0,-1.0},{1.0,-1.0},{1.0,1.0},{-1.0,1.0}};
//     _domain.setVertexList(p);
//     vector<int> _arg({5,5});
//     Mesh<2> mesh(_domain,_arg);//n等分
//     _mesh = mesh;
// }
void Possion_2D::setup_system()
{
    /*Enumerating is done by using DistributeDofs(),2-order element*/
    _mesh->distributeDofs();/*<  Q1网格对应Q1单元，Q2网格对应Q2单元.关注成员变量是否对应 */
    std::cout << "Number of degrees of freedom: " << _mesh->getTotalNumDofs() << std::endl;/*<  changed with order,same for below,same for getBoundaryDofs(2)  */
    _element->setMesh(*_mesh);/*<  right!! */

    /**
     * set the sizes of system_matrix, right hand side vector ans the solution vector;
     */
    system_matrix = Eigen::SparseMatrix<double>(int(_mesh->getTotalNumDofs()),int(_mesh->getTotalNumDofs()));
    system_matrix.setZero();
    system_rhs = Eigen::VectorXd::Zero(int(_mesh->getTotalNumDofs()));
    solution = Eigen::VectorXd::Zero(int(_mesh->getTotalNumDofs()));

}

void Possion_2D::assemble_system()
{
    /**
     * Get the right_hand_side function and boundary function
     */
    auto right_hand_side = _eqation->RightHandsTermFunction();
    // std::function<double(Dofs<2U>&)> boundary_function = BoundaryFunction;
 
    /**
     * for each Grid, we construct a cell_matirx and cell_rhs.
     */
    const unsigned int dofs_per_cell = _element->n_Dofs();
    Eigen::MatrixXd cell_matrix(dofs_per_cell,dofs_per_cell);
    Eigen::VectorXd cell_rhs(dofs_per_cell);

    /*assemble the global matrix and vector cell-by-cell.*/
    for (auto &cell : _mesh->getGridVector())
    {
    	// bind the TemplateElement with current cell.
	    _element->reinit(cell);
	    cell_matrix = Eigen::MatrixXd::Zero(dofs_per_cell,dofs_per_cell);
	    cell_rhs    = Eigen::VectorXd::Zero(dofs_per_cell);
	    for (int q_index = 0; q_index < _element->n_GaussPnt();++q_index)
        {
	        Point<2> pnt = _element->GaussionPoint(q_index);
	        for (int i = 0; i < dofs_per_cell; ++i)
	        {
		        for (int j = 0; j < dofs_per_cell; ++j)
		        {
		            double InnerProduct = 0;
		            for(int k = 0;k < _element->gradient(pnt[0],pnt[1],i+1).size();k++)
		            {
		                InnerProduct += _element->gradient(pnt[0],pnt[1],i+1)[k] *
			            _element->gradient(pnt[0],pnt[1],j+1)[k];
		            }
		            cell_matrix(i, j) +=(InnerProduct *
				    _element->det_Jacobi(pnt[0], pnt[1]) *
			        _element->GaussionWeight(q_index));
		        }
	        }
	        double xi = pnt[0],eta = pnt[1];
	        double x = _element->Global_x(xi, eta);
	        double y = _element->Global_y(xi, eta);
	        double p[2] = {x,y};
	        for (int i = 0; i < dofs_per_cell; ++i)
	        {
		        // Point<2> p({x,y});
		        // Dofs<2> globalDof(p);
		        cell_rhs(i) += ( _element->phi(xi, eta, i + 1) *
			    right_hand_side(p) *
			    _element->det_Jacobi(xi, eta) *
			    _element->GaussionWeight(q_index));
	        }
	    }
        for (int i = 0; i < dofs_per_cell; ++i)
	    {
	        for (int j = 0; j < dofs_per_cell; ++j)
	        {
		        int row = _element->GetGlobalIndex(i+1);
		        int column = _element->GetGlobalIndex(j+1);
		        system_matrix.coeffRef(row, column)+=cell_matrix(i, j);
	        }
	    }
	    for (int i = 0; i < dofs_per_cell; ++i)
	        system_rhs[_element->GetGlobalIndex(i+1)] += cell_rhs(i);
    }
    // cout << Eigen::MatrixXd(system_matrix) << endl;
    system_matrix.makeCompressed();
}

void Possion_2D::dealwithBoundaryCondition()
{
    std::vector<Dofs<2>> boundary_dofs = _mesh->getBoundaryDofs();
    std::map<int, double> boundary_values;
    auto boundary_function = _eqation->BoundaryConditionFunction();
    boundary_values = interpolateBoundaryValues(boundary_function,boundary_dofs);
    applyBoundaryValues(boundary_values,system_matrix,system_rhs);
}

void Possion_2D::solve()
{
    /*<  CG solve */
    Eigen::ConjugateGradient<Eigen::SparseMatrix<double> > Solver_sparse;
    Solver_sparse.setTolerance(1e-12);
    Solver_sparse.compute(system_matrix);
    solution = Solver_sparse.solve(system_rhs);
    /*<  AMG solve */
    // AMG_Multigrid AmgSolve(system_matrix, system_rhs, 1e-12, GS, 2.0 / 3, 0.5);
    // AmgSolve.Solve();
    // solution = AmgSolve.GetResult();
}

SpMat Possion_2D::A() const
{
    return system_matrix;
}
RHS Possion_2D::Rhs() const
{
    return system_rhs;
}

/**
 * @brief 返回map，存放边界上自由度的全局标号，和边界函数(仅Dirichlet)在对应自由度上的函数值。
 *
 * @param _boundary_function 边界函数
 * @param _boundary_dofs 边界上的自由度
 */
template <unsigned int DIM>
std::map<int,double> interpolateBoundaryValues(auto& _boundary_function, vector<Dofs<DIM>>& _boundary_dofs){
    std::map<int,double> boundary_values;
    for(Dofs<DIM>& dofs : _boundary_dofs)
    {
	int index = dofs.getGlobalIndex();
	double p[DIM];
	for(int i = 0;i < DIM;i++){
	    p[i] = dofs[i];
	}
    double value;
    if(p[1] == -1)
	    value = _boundary_function[0](p);
    else if(p[0] == 1)
	    value = _boundary_function[1](p);
    else if(p[1] == 1)
	    value = _boundary_function[2](p);
    else
	    value = _boundary_function[3](p);
	boundary_values.insert({index,value});
    }
    return boundary_values;
}
/**
 * @brief 对刚度矩阵和右端项应用Dirichlet条件
 * 
 * @param system_matrix 刚度矩阵
 * @param system_rhs 右端项
 */

void applyBoundaryValues(std::map<int,double>& boundary_values,Eigen::SparseMatrix<double>& system_matrix,Eigen::VectorXd& system_rhs){
    for(auto& it : boundary_values)
    {
        //边界自由度的全局编号。
	int k = it.first; 
        //对应自由的边界函数值。
	double bnd_value = it.second;
	system_rhs[k] = bnd_value * system_matrix.coeffRef(k,k);
	for(Eigen::SparseMatrix<double>::InnerIterator it(system_matrix,k);it;++it)
	{
	    int row = it.row();
	    if(row == k)
		continue;
	    system_rhs[row] -= system_matrix.coeffRef(k,row)* bnd_value;
	    system_matrix.coeffRef(k,row) = 0.0;
	    system_matrix.coeffRef(row,k) = 0.0;
	}

    }
}
#else
#endif
